"""
智谱AI向量模型实现
"""

from typing import List

from zhipuai import ZhipuAI

from config.logging_config import get_logger
from .BaseEmb import BaseEmb, ConfigurationError, EmbeddingError

logger = get_logger(__name__)


class ZhiPuEmb(BaseEmb):
    """智谱AI向量模型实现"""

    # 支持的模型列表
    SUPPORTED_MODELS = [
        "embedding-2",
        "embedding-3",
        "text-embedding-3-large",
        "text-embedding-3-small"
    ]

    def __init__(self, model_name: str, api_key: str, **kwargs):
        """
        初始化智谱AI向量模型
        
        Args:
            model_name: 模型名称
            api_key: API密钥
            **kwargs: 其他配置参数
        """
        self.api_key = api_key
        super().__init__(model_name, **kwargs)

        # 初始化客户端
        try:
            self.client = ZhipuAI(api_key=api_key)
        except Exception as e:
            raise ConfigurationError(f"初始化智谱AI客户端失败: {str(e)}")

    def _validate_config(self) -> None:
        """验证配置参数"""
        if not self.api_key:
            raise ConfigurationError("智谱AI API密钥不能为空")

        if self.model_name not in self.SUPPORTED_MODELS:
            logger.warning(f"模型 {self.model_name} 不在支持列表中，可能无法正常工作")

    def get_emb(self, text: str) -> List[float]:
        """
        获取单个文本的向量表示
        
        Args:
            text: 输入文本
            
        Returns:
            向量列表
            
        Raises:
            EmbeddingError: 向量化失败时抛出
        """
        if not text or not text.strip():
            raise ValueError("输入文本不能为空")

        try:
            response = self.client.embeddings.create(
                model=self.model_name,
                input=text.strip(),
            )

            if not response.data or len(response.data) == 0:
                raise EmbeddingError("API返回空结果")

            return response.data[0].embedding

        except Exception as e:
            if isinstance(e, (ValueError, EmbeddingError)):
                raise
            logger.error(f"智谱AI向量化失败: {str(e)}")
            raise EmbeddingError(f"向量化失败: {str(e)}")

    def encode(self, texts: List[str]) -> List[List[float]]:
        """
        编码多个文本
        
        Args:
            texts: 文本列表
            
        Returns:
            向量列表的列表
        """
        if not texts:
            return []

        # 过滤空文本
        valid_texts = [text.strip() for text in texts if text and text.strip()]
        if not valid_texts:
            raise ValueError("没有有效的输入文本")

        try:
            # 智谱AI支持批量处理
            response = self.client.embeddings.create(
                model=self.model_name,
                input=valid_texts,
            )

            if not response.data or len(response.data) != len(valid_texts):
                raise EmbeddingError("API返回结果数量不匹配")

            # 按顺序返回结果
            results = []
            for item in response.data:
                results.append(item.embedding)

            return results

        except Exception as e:
            if isinstance(e, (ValueError, EmbeddingError)):
                raise
            logger.error(f"智谱AI批量向量化失败: {str(e)}")
            raise EmbeddingError(f"批量向量化失败: {str(e)}")

    def encode_batch(self, texts: List[str], batch_size: int = 100) -> List[List[float]]:
        """
        批量编码文本（智谱AI优化版本）
        
        Args:
            texts: 文本列表
            batch_size: 批次大小，智谱AI建议不超过100
            
        Returns:
            向量列表的列表
        """
        if not texts:
            return []

        # 智谱AI建议批次大小不超过100
        batch_size = min(batch_size, 100)

        if len(texts) <= batch_size:
            return self.encode(texts)

        results = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            try:
                batch_results = self.encode(batch)
                results.extend(batch_results)
            except Exception as e:
                logger.error(f"批次 {i // batch_size + 1} 处理失败: {str(e)}")
                # 可以选择跳过失败的批次或者抛出异常
                raise

        return results

    def get_model_info(self) -> dict:
        """
        获取模型信息
        
        Returns:
            模型信息字典
        """
        return {
            "provider": "ZhiPu AI",
            "model_name": self.model_name,
            "supported_models": self.SUPPORTED_MODELS,
            "max_batch_size": 100,
            "dimension": self.get_dimension() if hasattr(self, '_dimension_cache') else None
        }

    def test_connection(self) -> bool:
        """
        测试连接是否正常
        
        Returns:
            连接是否正常
        """
        try:
            test_vector = self.get_emb("测试连接")
            return len(test_vector) > 0
        except Exception as e:
            logger.error(f"智谱AI连接测试失败: {str(e)}")
            return False
